from flask import Blueprint, request, jsonify
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from langchain_ollama.embeddings import OllamaEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from env import (
    model_url,
    model_name,
    collection_name,
    mongodb_uri,
    db_name,
)
from .utils import init_chroma, extract_list, build_list_to_json
from pymongo import MongoClient
import os

build_chroma_bp = Blueprint("build_chroma", __name__)

# 初始化嵌入模型
embedding_model = OllamaEmbeddings(base_url=model_url, model=model_name)

print(f"collection_name: {collection_name}")


# 输出文件
def build_data():
    client = MongoClient(mongodb_uri)

    db = client[db_name]
    markdowns_collection = db["markdowns"]
    markdowns_data = markdowns_collection.find()

    def get_content(item):
        return item["content"]

    return [
        {
            "page_content": get_content(markdown),
            "metadata": {"title": markdown["title"]}
            | build_list_to_json(extract_list(get_content(markdown))),
        }
        for markdown in markdowns_data
    ]


chroma_collection = init_chroma(embedding_model)


@build_chroma_bp.route("/build_chroma", methods=["POST"])
def build_chroma():
    try:
        print("build_chroma")
        data = build_data()

        # 创建文档列表
        documents = [
            Document(page_content=item["page_content"], metadata=item["metadata"])
            for item in data
        ]

        # 提取文档的嵌入向量
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
        split_texts = text_splitter.split_text(
            " ".join([doc.page_content for doc in documents])
        )
        embeddings = [
            embedding_model.embed_documents([text])[0] for text in split_texts
        ]

        # 准备数据点
        data_points = []
        for i, (text, embedding) in enumerate(zip(split_texts, embeddings)):
            data_points.append((f"vector{i}", embedding, {"text": text}, text))

        # 添加数据到 Chroma 集合
        ids = [dp[0] for dp in data_points]
        embeddings = [dp[1] for dp in data_points]
        metadatas = [dp[2] for dp in data_points]
        texts = [dp[3] for dp in data_points]

        chroma_collection.add_texts(
            texts, ids=ids, embeddings=embeddings, metadatas=metadatas
        )

        return jsonify({"message": "Data added successfully."}), 200
    except Exception as e:
        print(f"encounter error:{e}")
        return jsonify({"error": str(e)}), 500
